{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "792d95f9",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/tensorrt_torchtrt_efficientnet/nvidia_logo.png\" width=\"90px\">\n",
    "\n",
    "# PySpark PyTorch Inference\n",
    "\n",
    "### Regression\n",
    "\n",
    "In this notebook, we will train an MLP to perform regression on the California housing dataset, and load it for distributed inference with Spark.  \n",
    "\n",
    "Based on: https://github.com/christianversloot/machine-learning-articles/blob/main/how-to-create-a-neural-network-for-regression-with-pytorch.md  \n",
    "\n",
    "We also demonstrate accelerated inference via Torch-TensorRT model compilation.   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "75930360-c5ce-49ef-a69a-da88fa69a2ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import os\n",
    "import shutil\n",
    "import numpy as np\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "from sklearn.datasets import fetch_california_housing\n",
    "from sklearn.preprocessing import StandardScaler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1de685f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.mkdir('models') if not os.path.exists('models') else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6d5bc0c7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.5.1+cu124'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8754b174",
   "metadata": {},
   "source": [
    "### Load Dataset\n",
    "\n",
    "Each label corresponds to the average house value in units of 100,000, which we'll try to predict using the following features:  \n",
    "['MedInc', 'HouseAge', 'AveRooms', 'AveBedrms', 'Population', 'AveOccup', 'Latitude', 'Longitude']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2bee64cf-a44a-4aff-82db-c64ee3a8b0e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "X, y = fetch_california_housing(return_X_y=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8644e508-5e4c-4cdd-9ed1-9235887d9659",
   "metadata": {},
   "outputs": [],
   "source": [
    "class HousingDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, X, y, scale_data=True):\n",
    "        if not torch.is_tensor(X) and not torch.is_tensor(y):\n",
    "            # Apply scaling if necessary\n",
    "            if scale_data:\n",
    "                X = StandardScaler().fit_transform(X)\n",
    "            self.X = torch.from_numpy(X.astype(np.float32))\n",
    "            self.y = torch.from_numpy(y.astype(np.float32))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.X)\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        return self.X[i], self.y[i]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cc6b55c3-dc7b-4831-9943-83efd48091bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = HousingDataset(X, y)\n",
    "trainloader = torch.utils.data.DataLoader(\n",
    "    dataset, batch_size=10, shuffle=True, num_workers=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d868f39d-4695-4110-91d2-6f7a09d73b93",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([[ 6.5799e-01,  4.2594e-01, -1.4755e-01, -2.3638e-01, -4.0221e-01,\n",
       "          -5.6793e-02,  8.8868e-01, -1.3528e+00],\n",
       "         [ 6.7288e-01, -1.0043e+00,  5.7486e-01, -1.6537e-01, -3.3422e-01,\n",
       "          -6.4971e-02, -1.2790e+00,  1.2327e+00],\n",
       "         [-1.1616e-01,  2.8646e-02, -1.7830e-01, -2.3817e-01, -6.7154e-01,\n",
       "          -3.6429e-02, -1.3258e+00,  1.2726e+00],\n",
       "         [-3.2513e-01, -6.8648e-01, -3.4226e-01, -8.2805e-02,  5.1239e+00,\n",
       "           2.6689e-02, -7.7338e-01,  8.3340e-01],\n",
       "         [ 1.0892e-01, -1.2427e+00,  2.7819e-01, -8.7150e-02,  3.0158e-01,\n",
       "          -1.8564e-02, -1.1245e+00,  1.1628e+00],\n",
       "         [-8.6416e-02,  5.8485e-01, -7.8085e-02,  8.1655e-02, -6.7154e-01,\n",
       "          -1.6053e-02, -3.4733e-01,  1.2577e+00],\n",
       "         [-1.2463e-01,  1.0810e-01,  2.6662e-01, -1.0883e-01,  3.4839e-01,\n",
       "          -2.3125e-02, -7.7338e-01,  1.3325e+00],\n",
       "         [-9.2662e-01, -1.6400e+00, -2.4824e-01,  6.0041e-01,  6.3361e-01,\n",
       "          -1.0926e-01, -8.8574e-01,  1.2826e+00],\n",
       "         [ 2.0038e+00, -6.0702e-01,  8.4770e-01, -2.1254e-01,  1.3745e+00,\n",
       "          -5.0489e-03, -6.5165e-01,  2.5441e-01],\n",
       "         [-3.9250e-01,  1.0616e+00, -1.8614e-01, -1.7073e-01, -3.8543e-01,\n",
       "          -8.1186e-02,  1.0806e+00, -1.3827e+00]]),\n",
       " tensor([3.1090, 1.8430, 1.6890, 1.8670, 1.9600, 0.6200, 0.9860, 0.9440, 3.9120,\n",
       "         1.4390])]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "next(iter(trainloader))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e817b9a",
   "metadata": {},
   "source": [
    "### Create and Train Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9a441b60-dca4-44d2-bc1c-aa7336d704bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.layers = nn.Sequential(\n",
    "            nn.Linear(8, 64),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(64, 32),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(32, 1)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.layers(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15cff2b4-9d23-4d2b-808a-a5edb8eda135",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using cuda device\n"
     ]
    }
   ],
   "source": [
    "# Initialize the MLP\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f\"Using {device} device\")\n",
    "mlp = MLP().to(device)\n",
    "\n",
    "# Define the loss function and optimizer\n",
    "loss_function = nn.L1Loss()\n",
    "optimizer = torch.optim.Adam(mlp.parameters(), lr=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5e2db3f9-5db8-4b42-89ad-e77f23c4c1fe",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting epoch 1\n",
      "Loss after mini-batch     1: 0.004\n",
      "Loss after mini-batch   201: 0.701\n",
      "Loss after mini-batch   401: 0.463\n",
      "Loss after mini-batch   601: 0.329\n",
      "Loss after mini-batch   801: 0.285\n",
      "Loss after mini-batch  1001: 0.253\n",
      "Loss after mini-batch  1201: 0.247\n",
      "Loss after mini-batch  1401: 0.234\n",
      "Loss after mini-batch  1601: 0.232\n",
      "Loss after mini-batch  1801: 0.217\n",
      "Loss after mini-batch  2001: 0.211\n",
      "Starting epoch 2\n",
      "Loss after mini-batch     1: 0.001\n",
      "Loss after mini-batch   201: 0.205\n",
      "Loss after mini-batch   401: 0.212\n",
      "Loss after mini-batch   601: 0.206\n",
      "Loss after mini-batch   801: 0.205\n",
      "Loss after mini-batch  1001: 0.202\n",
      "Loss after mini-batch  1201: 0.202\n",
      "Loss after mini-batch  1401: 0.204\n",
      "Loss after mini-batch  1601: 0.198\n",
      "Loss after mini-batch  1801: 0.188\n",
      "Loss after mini-batch  2001: 0.188\n",
      "Starting epoch 3\n",
      "Loss after mini-batch     1: 0.001\n",
      "Loss after mini-batch   201: 0.197\n",
      "Loss after mini-batch   401: 0.193\n",
      "Loss after mini-batch   601: 0.196\n",
      "Loss after mini-batch   801: 0.189\n",
      "Loss after mini-batch  1001: 0.183\n",
      "Loss after mini-batch  1201: 0.191\n",
      "Loss after mini-batch  1401: 0.193\n",
      "Loss after mini-batch  1601: 0.181\n",
      "Loss after mini-batch  1801: 0.185\n",
      "Loss after mini-batch  2001: 0.181\n",
      "Starting epoch 4\n",
      "Loss after mini-batch     1: 0.001\n",
      "Loss after mini-batch   201: 0.190\n",
      "Loss after mini-batch   401: 0.181\n",
      "Loss after mini-batch   601: 0.189\n",
      "Loss after mini-batch   801: 0.180\n",
      "Loss after mini-batch  1001: 0.184\n",
      "Loss after mini-batch  1201: 0.180\n",
      "Loss after mini-batch  1401: 0.180\n",
      "Loss after mini-batch  1601: 0.184\n",
      "Loss after mini-batch  1801: 0.186\n",
      "Loss after mini-batch  2001: 0.179\n",
      "Starting epoch 5\n",
      "Loss after mini-batch     1: 0.000\n",
      "Loss after mini-batch   201: 0.181\n",
      "Loss after mini-batch   401: 0.177\n",
      "Loss after mini-batch   601: 0.185\n",
      "Loss after mini-batch   801: 0.179\n",
      "Loss after mini-batch  1001: 0.178\n",
      "Loss after mini-batch  1201: 0.173\n",
      "Loss after mini-batch  1401: 0.185\n",
      "Loss after mini-batch  1601: 0.177\n",
      "Loss after mini-batch  1801: 0.181\n",
      "Loss after mini-batch  2001: 0.178\n",
      "Training process has finished.\n"
     ]
    }
   ],
   "source": [
    "# Run the training loop\n",
    "for epoch in range(0, 5):  # 5 epochs at maximum\n",
    "\n",
    "    # Print epoch\n",
    "    print(f'Starting epoch {epoch+1}')\n",
    "\n",
    "    # Set current loss value\n",
    "    current_loss = 0.0\n",
    "\n",
    "    # Iterate over the DataLoader for training data\n",
    "    for i, data in enumerate(trainloader, 0):\n",
    "\n",
    "        # Get and prepare inputs\n",
    "        inputs, targets = data\n",
    "        inputs, targets = inputs.to(device), targets.to(device)\n",
    "        targets = targets.reshape((targets.shape[0], 1))\n",
    "\n",
    "        # Zero the gradients\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # Perform forward pass\n",
    "        outputs = mlp(inputs)\n",
    "\n",
    "        # Compute loss\n",
    "        loss = loss_function(outputs, targets)\n",
    "\n",
    "        # Perform backward pass\n",
    "        loss.backward()\n",
    "\n",
    "        # Perform optimization\n",
    "        optimizer.step()\n",
    "\n",
    "        # Print statistics\n",
    "        current_loss += loss.item()\n",
    "        if i % 200 == 0:\n",
    "            print('Loss after mini-batch %5d: %.3f' %\n",
    "                  (i + 1, current_loss / 500))\n",
    "            current_loss = 0.0\n",
    "\n",
    "# Process is complete.\n",
    "print('Training process has finished.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "352539f5",
   "metadata": {},
   "source": [
    "### Save Model State Dict\n",
    "This saves the serialized object to disk using pickle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b950a3ed-ffe1-477f-a84f-f71c85dbf9ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved PyTorch Model State to models/housing_model.pt\n"
     ]
    }
   ],
   "source": [
    "torch.save(mlp.state_dict(), \"models/housing_model.pt\")\n",
    "print(\"Saved PyTorch Model State to models/housing_model.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0060fcca",
   "metadata": {},
   "source": [
    "### Save Model as TorchScript\n",
    "This saves an [intermediate representation of the compute graph](https://pytorch.org/tutorials/beginner/saving_loading_models.html#export-load-model-in-torchscript-format), which does not require pickle (or even python). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20fedb5d-c59e-4b0b-ba91-3dd15df1f09e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved TorchScript Model to models/ts_housing_model.pt\n"
     ]
    }
   ],
   "source": [
    "scripted = torch.jit.script(mlp)\n",
    "scripted.save(\"models/ts_housing_model.pt\")\n",
    "print(\"Saved TorchScript Model to models/ts_housing_model.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3101c0fe-65f1-411e-9192-e8a6b585ba0d",
   "metadata": {},
   "source": [
    "### Load and Test from Model State"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "7411b00f-88d2-40f5-b716-a26733c968ff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loaded_mlp = MLP().to(device)\n",
    "loaded_mlp.load_state_dict(torch.load(\"models/housing_model.pt\", weights_only=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "e226f449-2931-4492-9003-503cdc61f061",
   "metadata": {},
   "outputs": [],
   "source": [
    "testX, testY = next(iter(trainloader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d46af47e-db7e-42ee-9bd3-6e7d93850be3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predictions:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[2.3652],\n",
       "        [1.8444],\n",
       "        [2.4587],\n",
       "        [3.1243],\n",
       "        [2.2726],\n",
       "        [2.1818],\n",
       "        [1.5222],\n",
       "        [0.5554],\n",
       "        [2.2508],\n",
       "        [3.5971]], device='cuda:0', grad_fn=<AddmmBackward0>)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"Predictions:\")\n",
    "loaded_mlp(testX.to(device))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "13ae2c0f-1da5-45a4-bf32-ed8b562d7907",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Labels:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([2.7370, 2.2110, 2.5360, 2.6330, 1.6540, 2.3360, 1.4600, 0.6590, 2.6380,\n",
       "        3.6220])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"Labels:\")\n",
    "testY"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bcd329d",
   "metadata": {},
   "source": [
    "### Load and Test from TorchScript"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "422e317f-c9bd-4f76-9463-7af2935d401d",
   "metadata": {},
   "outputs": [],
   "source": [
    "scripted_mlp = torch.jit.load(\"models/ts_housing_model.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cda8ec8-644e-4888-bfa0-b79425ece7c3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predictions:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([2.3652, 1.8444, 2.4587, 3.1243, 2.2726, 2.1818, 1.5222, 0.5554, 2.2508,\n",
       "        3.5971], device='cuda:0', grad_fn=<ViewBackward0>)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"Predictions:\")\n",
    "scripted_mlp(testX.to(device)).flatten()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a3b64e4",
   "metadata": {},
   "source": [
    "### Compile using the Torch JIT Compiler\n",
    "This leverages the [Torch-TensorRT inference compiler](https://pytorch.org/TensorRT/) for accelerated inference on GPUs using the `torch.compile` JIT interface under the hood. The compiler stack returns a [boxed-function](http://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/) that triggers compilation on the first call.  \n",
    "\n",
    "Modules compiled in this fashion are [not serializable with pickle](https://github.com/pytorch/pytorch/issues/101107#issuecomment-1542688089), so we cannot send the compiled model directly to Spark.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c613f24e",
   "metadata": {},
   "source": [
    "(You may see a warning about modelopt quantization. This is safe to ignore, as [implicit quantization](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#intro-quantization) is deprecated in the latest TensorRT. See [this link](https://pytorch.org/TensorRT/tutorials/_rendered_examples/dynamo/vgg16_fp8_ptq.html) for a guide to explicit quantization.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "9ffb27fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch_tensorrt as trt\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "e0c10f90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional: set the filename for the TensorRT timing cache\n",
    "timestamp = time.time()\n",
    "timing_cache = f\"/tmp/timing_cache-{timestamp}.bin\"\n",
    "with open(timing_cache, \"wb\") as f:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "b4aa2523",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs_bs1 = torch.randn((10, 8), dtype=torch.float).to(\"cuda\")\n",
    "# This indicates dimension 0 of inputs_bs1 is dynamic with a range of values [1, 50]. No recompilation will happen when the batch size changes.\n",
    "torch._dynamo.mark_dynamic(inputs_bs1, 0, min=1, max=50)\n",
    "trt_model = trt.compile(\n",
    "    loaded_mlp,\n",
    "    ir=\"torch_compile\",\n",
    "    inputs=inputs_bs1,\n",
    "    enabled_precisions={torch.float},\n",
    "    timing_cache_path=timing_cache,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "a5da8cab",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:torch_tensorrt.dynamo._compiler:Node linear_default of op type call_function does not have metadata. This could sometimes lead to undefined behavior.\n",
      "WARNING:torch_tensorrt.dynamo._compiler:Some nodes do not have metadata (shape and dtype information). This could lead to problems sometimes if the graph has PyTorch and TensorRT segments.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predictions:\n",
      "tensor([[2.3652],\n",
      "        [1.8444],\n",
      "        [2.4587],\n",
      "        [3.1243],\n",
      "        [2.2726],\n",
      "        [2.1818],\n",
      "        [1.5222],\n",
      "        [0.5554],\n",
      "        [2.2508],\n",
      "        [3.5971]], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "stream = torch.cuda.Stream()\n",
    "with torch.no_grad(), torch.cuda.stream(stream):\n",
    "    testX = testX.to(device)\n",
    "    print(\"Predictions:\")\n",
    "    print(trt_model(testX))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2c55e07",
   "metadata": {},
   "source": [
    "### Compile using the Torch-TensorRT AOT Compiler\n",
    "Alternatively, use the Torch-TensorRT Dynamo backend for Ahead-of-Time (AOT) compilation to eagerly optimize the model in an explicit compilation phase. We first export the model to produce a traced graph representing the Tensor computation in an AOT fashion, which produces a `ExportedProgram` object which can be [serialized and reloaded](https://pytorch.org/TensorRT/user_guide/saving_models.html). We can then compile this IR using the Torch-TensorRT AOT compiler for inference.   \n",
    "\n",
    "[Read the docs](https://pytorch.org/TensorRT/user_guide/torch_tensorrt_explained.html) for more information on JIT vs AOT compilation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "bf36a50d",
   "metadata": {},
   "outputs": [],
   "source": [
    "example_inputs = (torch.randn((10, 8), dtype=torch.float).to(\"cuda\"),)\n",
    "\n",
    "# Mark dim 1 (batch size) as dynamic\n",
    "batch = torch.export.Dim(\"batch\", min=1, max=64)\n",
    "# Produce traced graph in ExportedProgram format\n",
    "exp_program = torch.export.export(loaded_mlp, args=example_inputs, dynamic_shapes={\"x\": {0: batch}})\n",
    "# Compile the traced graph to produce an optimized module\n",
    "trt_gm = trt.dynamo.compile(exp_program,\n",
    "                            tuple(example_inputs),\n",
    "                            enabled_precisions={torch.float},\n",
    "                            timing_cache_path=timing_cache,\n",
    "                            workspace_size=1<<30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "4fc4efd5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'torch.export.exported_program.ExportedProgram'>\n",
      "<class 'torch.fx.graph_module.GraphModule.__new__.<locals>.GraphModuleImpl'>\n"
     ]
    }
   ],
   "source": [
    "print(type(exp_program))\n",
    "print(type(trt_gm))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "1bcf9c47",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predictions:\n",
      "tensor([[2.3653],\n",
      "        [1.8443],\n",
      "        [2.4586],\n",
      "        [3.1242],\n",
      "        [2.2725],\n",
      "        [2.1815],\n",
      "        [1.5221],\n",
      "        [0.5556],\n",
      "        [2.2508],\n",
      "        [3.5971]], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "stream = torch.cuda.Stream()\n",
    "with torch.no_grad(), torch.cuda.stream(stream):\n",
    "    print(\"Predictions:\")\n",
    "    testX = testX.to(device)\n",
    "    print(trt_gm(testX))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0eeb957a",
   "metadata": {},
   "source": [
    "We can run the optimized module with a few different batch sizes (without recompilation!):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49f72c14",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output shapes:\n",
      "torch.Size([10, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([50, 1])\n"
     ]
    }
   ],
   "source": [
    "inputs = (torch.randn((10, 8), dtype=torch.float).cuda(),)\n",
    "inputs_bs1 = (torch.randn((1, 8), dtype=torch.float).cuda(),)\n",
    "inputs_bs50 = (torch.randn((50, 8), dtype=torch.float).cuda(),)\n",
    "\n",
    "stream = torch.cuda.Stream()\n",
    "with torch.no_grad(), torch.cuda.stream(stream):\n",
    "    print(\"Output shapes:\")\n",
    "    print(trt_gm(*inputs).shape)\n",
    "    print(trt_gm(*inputs_bs1).shape)\n",
    "    print(trt_gm(*inputs_bs50).shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4fef57d",
   "metadata": {},
   "source": [
    "We can serialize the ExportedProgram (a traced graph representing the model's forward function) using `torch.export.save` to be recompiled at a later date."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "876fea4a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved ExportedProgram to models/trt_housing_model.ep\n"
     ]
    }
   ],
   "source": [
    "torch.export.save(exp_program, \"models/trt_housing_model.ep\")\n",
    "print(\"Saved ExportedProgram to models/trt_housing_model.ep\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13631d1f-2c71-4bee-afcb-bd3b55ec87c5",
   "metadata": {},
   "source": [
    "## PySpark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb71dd36",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, struct, pandas_udf, array\n",
    "from pyspark.ml.functions import predict_batch_udf\n",
    "from pyspark.sql.types import *\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark import SparkConf\n",
    "import json\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6769c060",
   "metadata": {},
   "source": [
    "Check the cluster environment to handle any platform-specific Spark configurations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "f7727b58",
   "metadata": {},
   "outputs": [],
   "source": [
    "on_databricks = os.environ.get(\"DATABRICKS_RUNTIME_VERSION\", False)\n",
    "on_dataproc = os.environ.get(\"DATAPROC_IMAGE_VERSION\", False)\n",
    "on_standalone = not (on_databricks or on_dataproc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3b7d360",
   "metadata": {},
   "source": [
    "#### Create Spark Session\n",
    "\n",
    "For local standalone clusters, we'll connect to the cluster and create the Spark Session.  \n",
    "For CSP environments, Spark will either be preconfigured (Databricks) or we'll need to create the Spark Session (Dataproc)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "52e9dbdb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "25/02/04 13:46:28 WARN Utils: Your hostname, cb4ae00-lcedt resolves to a loopback address: 127.0.1.1; using 10.110.47.100 instead (on interface eno1)\n",
      "25/02/04 13:46:28 WARN Utils: Set SPARK_LOCAL_IP if you need to bind to another address\n",
      "Setting default log level to \"WARN\".\n",
      "To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).\n",
      "25/02/04 13:46:28 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable\n"
     ]
    }
   ],
   "source": [
    "conf = SparkConf()\n",
    "\n",
    "if 'spark' not in globals():\n",
    "    if on_standalone:\n",
    "        import socket\n",
    "        conda_env = os.environ.get(\"CONDA_PREFIX\")\n",
    "        hostname = socket.gethostname()\n",
    "        conf.setMaster(f\"spark://{hostname}:7077\")\n",
    "        conf.set(\"spark.pyspark.python\", f\"{conda_env}/bin/python\")\n",
    "        conf.set(\"spark.pyspark.driver.python\", f\"{conda_env}/bin/python\")\n",
    "\n",
    "    conf.set(\"spark.executor.cores\", \"8\")\n",
    "    conf.set(\"spark.task.resource.gpu.amount\", \"0.125\")\n",
    "    conf.set(\"spark.executor.resource.gpu.amount\", \"1\")\n",
    "    conf.set(\"spark.sql.execution.arrow.pyspark.enabled\", \"true\")\n",
    "    conf.set(\"spark.python.worker.reuse\", \"true\")\n",
    "    \n",
    "spark = SparkSession.builder.appName(\"spark-dl-examples\").config(conf=conf).getOrCreate()\n",
    "sc = spark.sparkContext"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3b9937e-2c70-4d67-b95f-4d9d5ab17c12",
   "metadata": {},
   "source": [
    "### Create Spark DataFrame from Pandas DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "cf35da14-61a3-4e7b-9d4f-086bf5e931b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "housing = fetch_california_housing()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "95148019-ea95-40e5-a529-fcdb9a06f928",
   "metadata": {},
   "outputs": [],
   "source": [
    "X = StandardScaler().fit_transform(housing.data.astype(np.float32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "f82d957c-6747-4408-aac8-45305afbfe5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "pdf = pd.DataFrame(X, columns=housing.feature_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "881afee9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------------+------------+-----------+------------+-----------+------------+----------+------------+\n",
      "|      MedInc|    HouseAge|   AveRooms|   AveBedrms| Population|    AveOccup|  Latitude|   Longitude|\n",
      "+------------+------------+-----------+------------+-----------+------------+----------+------------+\n",
      "|  0.20909257|  -1.1632254| 0.38946992|  0.04609274| -0.9806099| -0.07099328|0.61245227|-0.020113053|\n",
      "|-0.098627955|  0.34647804| 0.27216315|  -0.0129226| -0.6953838| -0.05380849| 1.0665938|  -1.2479742|\n",
      "| -0.66006273|   1.0616008|-0.55292207| -0.48945764|-0.13641118| 0.028952759| 1.1040496|  -1.3827378|\n",
      "|  0.08218294|   0.5848523|-0.13912922| -0.14707813|-0.19116047| -0.07136432|0.96827507|  -1.3028787|\n",
      "|   0.0784456|  -1.4810578| 0.57265776|  0.32067496|  1.0345173|-0.024157424| 1.4411427| -0.52423614|\n",
      "| -0.82318723| -0.36864465| 0.07829511|  -0.1808107|-0.67242444|-0.061470542| 1.9374212|  -1.0083897|\n",
      "|  0.59671736|   0.5848523| 0.19346413|  -0.1371872|-0.19645879| 0.009964322|0.96827507|  -1.2928978|\n",
      "|  -0.9612035|  -1.5605159|-0.56329846| 0.027148023|-0.71127874| -0.08471591| 0.5328614| -0.13990337|\n",
      "| -0.74344087|  -1.2426835| 0.27282518|   0.4037246| -0.9841421| -0.05610115| 1.2257773| -0.42940006|\n",
      "|   0.9784464|  -0.2891866| 0.24374022| -0.24670053| 0.28922042| -0.01102468| 1.1087307|  -1.2280084|\n",
      "|  -0.5070446|  -1.0043093|-0.78254056|0.0122275995|  2.8465424|-0.060435444| 0.8980464|  -1.2080427|\n",
      "| -0.18690155|   1.2205169|0.015323491|  0.12183313|-0.41015765|  0.04452552|  1.010412|  -1.3228445|\n",
      "|  -1.2551856|   1.6178073| -0.3341509|-0.060125165| -0.7554314| -0.08777025| 1.0291398|  -1.3477987|\n",
      "|   4.9607058|  -1.9578062|  1.4854684| -0.03948475|  2.1833694|0.0029250523|  1.024457|  -1.1581304|\n",
      "|  0.73652315|  -1.6399739|  0.7913185| -0.05238397|    1.67738|  0.01944797| 1.0993668|  -1.1331724|\n",
      "|   -0.505834|  0.18756187|-0.47093546| -0.24297306|-0.60619545| -0.10791535|  0.977639|  -1.2879055|\n",
      "| -0.88477343|-0.050812364| -0.6318951| -0.15244243| -0.5258376| -0.15618815| 0.9823201|  -1.2879055|\n",
      "| -0.42840376|   0.9821427| -0.2266495| -0.36083496| -0.6883194| -0.08552282| 0.5328614| -0.12493005|\n",
      "|   0.9369153|  -1.4810578|  0.6722208|-0.121177554|  0.3996021|  0.01291408| 1.1040496|  -1.1082181|\n",
      "| -0.80702734| -0.92485124|-0.26602685|  -0.1560743|  1.4398388| -0.09314839|0.55627036| -0.09498342|\n",
      "+------------+------------+-----------+------------+-----------+------------+----------+------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "schema = StructType([\n",
    "    StructField(\"MedInc\",FloatType(),True),\n",
    "    StructField(\"HouseAge\",FloatType(),True),\n",
    "    StructField(\"AveRooms\",FloatType(),True),\n",
    "    StructField(\"AveBedrms\",FloatType(),True),\n",
    "    StructField(\"Population\",FloatType(),True),\n",
    "    StructField(\"AveOccup\",FloatType(),True),\n",
    "    StructField(\"Latitude\",FloatType(),True),\n",
    "    StructField(\"Longitude\",FloatType(),True)\n",
    "])\n",
    "\n",
    "df = spark.createDataFrame(pdf, schema=schema).repartition(8)\n",
    "df.show(truncate=12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "7b33d367-fbf9-4918-b755-5447125547c4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "StructType([StructField('MedInc', FloatType(), True), StructField('HouseAge', FloatType(), True), StructField('AveRooms', FloatType(), True), StructField('AveBedrms', FloatType(), True), StructField('Population', FloatType(), True), StructField('AveOccup', FloatType(), True), StructField('Latitude', FloatType(), True), StructField('Longitude', FloatType(), True)])"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "751bff7a-b687-4184-b3fa-b5f5b46ef5d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = \"spark-dl-datasets/california_housing\"\n",
    "if on_databricks:\n",
    "    dbutils.fs.mkdirs(\"/FileStore/spark-dl-datasets\")\n",
    "    data_path = \"dbfs:/FileStore/\" + data_path\n",
    "\n",
    "df.write.mode(\"overwrite\").parquet(data_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88c3cd75",
   "metadata": {},
   "source": [
    "## Inference using Spark DL API\n",
    "\n",
    "Distributed inference using the PySpark [predict_batch_udf](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.ml.functions.predict_batch_udf.html#pyspark.ml.functions.predict_batch_udf):\n",
    "\n",
    "- predict_batch_fn uses PyTorch APIs to load the model and return a predict function which operates on numpy arrays \n",
    "- predict_batch_udf will convert the Spark DataFrame columns into numpy input batches for the predict function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "1e40c266-24de-454d-a776-f3716ba50e90",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = spark.read.parquet(data_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "5b144c17",
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = df.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "3d608e2f-66a8-44b5-9cde-5f7837bf4247",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get absolute path to model\n",
    "model_path = \"{}/models/trt_housing_model.ep\".format(os.getcwd())\n",
    "\n",
    "# For cloud environments, copy the model to the distributed file system.\n",
    "if on_databricks:\n",
    "    dbutils.fs.mkdirs(\"/FileStore/spark-dl-models\")\n",
    "    dbfs_model_path = \"/dbfs/FileStore/spark-dl-models/trt_housing_model.ep\"\n",
    "    shutil.copy(model_path, dbfs_model_path)\n",
    "    model_path = dbfs_model_path\n",
    "elif on_dataproc:\n",
    "    # GCS is mounted at /mnt/gcs by the init script\n",
    "    models_dir = \"/mnt/gcs/spark-dl/models\"\n",
    "    os.mkdir(models_dir) if not os.path.exists(models_dir) else None\n",
    "    gcs_model_path = models_dir + \"/trt_housing_model.ep\"\n",
    "    shutil.copy(model_path, gcs_model_path)\n",
    "    model_path = gcs_model_path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fd143e7",
   "metadata": {},
   "source": [
    "For inference on Spark, we'll load the ExportedProgram and compile the model with the Torch-TensorRT AOT compiler and cache on the executor. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "fc400771",
   "metadata": {},
   "outputs": [],
   "source": [
    "# A resource warning may occur due to unclosed file descriptors used by TensorRT across multiple PySpark daemon processes.\n",
    "# These can be safely ignored as the resources will be cleaned up when the worker processes terminate.\n",
    "\n",
    "import warnings\n",
    "warnings.simplefilter(\"ignore\", ResourceWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "a2f45f5d-c941-4197-a274-1eec2af3fca4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_batch_fn():\n",
    "    import torch\n",
    "    import torch_tensorrt as trt\n",
    "\n",
    "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "    if device != \"cuda\":\n",
    "        raise ValueError(\"This function uses the TensorRT model which requires a GPU device\")\n",
    "\n",
    "    example_inputs = (torch.randn((50, 8), dtype=torch.float).to(\"cuda\"),)\n",
    "    exp_program = torch.export.load(model_path)\n",
    "    trt_gm = trt.dynamo.compile(exp_program,\n",
    "                            tuple(example_inputs),\n",
    "                            enabled_precisions={torch.float},\n",
    "                            timing_cache_path=timing_cache,\n",
    "                            workspace_size=1<<30)\n",
    "\n",
    "    print(\"Model compiled.\")\n",
    "    \n",
    "    def predict(inputs):\n",
    "        stream = torch.cuda.Stream()\n",
    "        with torch.no_grad(), torch.cuda.stream(stream), trt.logging.errors():\n",
    "            print(f\"Predict {inputs.shape}\")\n",
    "            torch_inputs = torch.from_numpy(inputs).to(device)\n",
    "            outputs = trt_gm(torch_inputs) # .flatten()\n",
    "            return outputs.detach().cpu().numpy()\n",
    "\n",
    "    return predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "220a00a4-e842-4f5d-a4b3-7693d09e2d31",
   "metadata": {},
   "outputs": [],
   "source": [
    "regress = predict_batch_udf(predict_batch_fn,\n",
    "                             return_type=FloatType(),\n",
    "                             input_tensor_shapes=[[8]],\n",
    "                             batch_size=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "0f3bf287-8ffc-4456-8772-e97c418d6aee",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 7:==============>                                            (2 + 6) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 30.4 ms, sys: 13.1 ms, total: 43.5 ms\n",
      "Wall time: 10.1 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = df.withColumn(\"preds\", regress(struct(*columns)))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "6cd23b71-296d-4ce7-b56c-567cc2eec79c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 31.6 ms, sys: 7.39 ms, total: 39 ms\n",
      "Wall time: 263 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = df.withColumn(\"preds\", regress(array(*columns)))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "75d16bd5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 28.7 ms, sys: 6.67 ms, total: 35.4 ms\n",
      "Wall time: 296 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = df.withColumn(\"preds\", regress(array(*columns)))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "764a40d8-25f7-425c-ba03-fe8c45f4b063",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------------+------------+-----------+------------+-----------+------------+----------+------------+---------+\n",
      "|      MedInc|    HouseAge|   AveRooms|   AveBedrms| Population|    AveOccup|  Latitude|   Longitude|    preds|\n",
      "+------------+------------+-----------+------------+-----------+------------+----------+------------+---------+\n",
      "|  0.20909257|  -1.1632254| 0.38946992|  0.04609274| -0.9806099| -0.07099328|0.61245227|-0.020113053|1.3746364|\n",
      "|-0.098627955|  0.34647804| 0.27216315|  -0.0129226| -0.6953838| -0.05380849| 1.0665938|  -1.2479742|1.8087528|\n",
      "| -0.66006273|   1.0616008|-0.55292207| -0.48945764|-0.13641118| 0.028952759| 1.1040496|  -1.3827378|1.4245079|\n",
      "|  0.08218294|   0.5848523|-0.13912922| -0.14707813|-0.19116047| -0.07136432|0.96827507|  -1.3028787|2.3895802|\n",
      "|   0.0784456|  -1.4810578| 0.57265776|  0.32067496|  1.0345173|-0.024157424| 1.4411427| -0.52423614|1.3616933|\n",
      "| -0.82318723| -0.36864465| 0.07829511|  -0.1808107|-0.67242444|-0.061470542| 1.9374212|  -1.0083897|0.7539238|\n",
      "|  0.59671736|   0.5848523| 0.19346413|  -0.1371872|-0.19645879| 0.009964322|0.96827507|  -1.2928978|2.6816423|\n",
      "|  -0.9612035|  -1.5605159|-0.56329846| 0.027148023|-0.71127874| -0.08471591| 0.5328614| -0.13990337|1.1731354|\n",
      "| -0.74344087|  -1.2426835| 0.27282518|   0.4037246| -0.9841421| -0.05610115| 1.2257773| -0.42940006|1.0198532|\n",
      "|   0.9784464|  -0.2891866| 0.24374022| -0.24670053| 0.28922042| -0.01102468| 1.1087307|  -1.2280084| 2.708211|\n",
      "|  -0.5070446|  -1.0043093|-0.78254056|0.0122275995|  2.8465424|-0.060435444| 0.8980464|  -1.2080427|2.0327075|\n",
      "| -0.18690155|   1.2205169|0.015323491|  0.12183313|-0.41015765|  0.04452552|  1.010412|  -1.3228445|1.9909104|\n",
      "|  -1.2551856|   1.6178073| -0.3341509|-0.060125165| -0.7554314| -0.08777025| 1.0291398|  -1.3477987|1.2702764|\n",
      "|   4.9607058|  -1.9578062|  1.4854684| -0.03948475|  2.1833694|0.0029250523|  1.024457|  -1.1581304| 5.975229|\n",
      "|  0.73652315|  -1.6399739|  0.7913185| -0.05238397|    1.67738|  0.01944797| 1.0993668|  -1.1331724|1.9309721|\n",
      "|   -0.505834|  0.18756187|-0.47093546| -0.24297306|-0.60619545| -0.10791535|  0.977639|  -1.2879055|1.7610806|\n",
      "| -0.88477343|-0.050812364| -0.6318951| -0.15244243| -0.5258376| -0.15618815| 0.9823201|  -1.2879055| 1.655031|\n",
      "| -0.42840376|   0.9821427| -0.2266495| -0.36083496| -0.6883194| -0.08552282| 0.5328614| -0.12493005|1.1175063|\n",
      "|   0.9369153|  -1.4810578|  0.6722208|-0.121177554|  0.3996021|  0.01291408| 1.1040496|  -1.1082181|2.1779811|\n",
      "| -0.80702734| -0.92485124|-0.26602685|  -0.1560743|  1.4398388| -0.09314839|0.55627036| -0.09498342|0.9102398|\n",
      "+------------+------------+-----------+------------+-----------+------------+----------+------------+---------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "preds.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "0aa85f81",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This will clear the engine cache (containing previously compiled TensorRT engines) and reset the CUDA Context.\n",
    "torch._dynamo.reset()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53536808",
   "metadata": {},
   "source": [
    "## Using Triton Inference Server\n",
    "In this section, we demonstrate integration with the [Triton Inference Server](https://developer.nvidia.com/nvidia-triton-inference-server), an open-source, GPU-accelerated serving solution for DL.  \n",
    "We use [PyTriton](https://github.com/triton-inference-server/pytriton), a Flask-like framework that handles client/server communication with the Triton server.  \n",
    "\n",
    "The process looks like this:\n",
    "- Distribute a PyTriton task across the Spark cluster, instructing each node to launch a Triton server process.\n",
    "- Define a Triton inference function, which contains a client that binds to the local server on a given node and sends inference requests.\n",
    "- Wrap the Triton inference function in a predict_batch_udf to launch parallel inference requests using Spark.\n",
    "- Finally, distribute a shutdown signal to terminate the Triton server processes on each node.\n",
    "\n",
    "<img src=\"../images/spark-server.png\" alt=\"drawing\" width=\"700\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "a9ab4cdf-8103-447e-9ac8-944e2e527239",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b77dc96",
   "metadata": {},
   "source": [
    "Import the helper class from server_utils.py:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "1ac83062",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.addPyFile(\"server_utils.py\")\n",
    "\n",
    "from server_utils import TritonServerManager"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4cc5d81",
   "metadata": {},
   "source": [
    "Define the Triton Server function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "6632636e-67a3-406c-832c-758aac4245fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def triton_server(ports, model_path):\n",
    "    import time\n",
    "    import signal\n",
    "    import numpy as np\n",
    "    import torch\n",
    "    import torch_tensorrt as trt\n",
    "    from pytriton.decorators import batch\n",
    "    from pytriton.model_config import DynamicBatcher, ModelConfig, Tensor\n",
    "    from pytriton.triton import Triton, TritonConfig\n",
    "    from pyspark import TaskContext\n",
    "\n",
    "    print(f\"SERVER: Initializing model on worker {TaskContext.get().partitionId()}.\")\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    \n",
    "    exp_program = torch.export.load(model_path)\n",
    "    example_inputs = (torch.randn((50, 8), dtype=torch.float).to(\"cuda\"),)\n",
    "    trt_gm = trt.dynamo.compile(exp_program,\n",
    "                            tuple(example_inputs),\n",
    "                            enabled_precisions={torch.float},\n",
    "                            workspace_size=1<<30)\n",
    "\n",
    "    print(\"SERVER: Compiled model.\")\n",
    "\n",
    "    @batch\n",
    "    def _infer_fn(**inputs):\n",
    "        features = inputs[\"features\"]\n",
    "        if len(inputs[\"features\"]) != 1:\n",
    "            features = np.squeeze(features)\n",
    "        stream = torch.cuda.Stream()\n",
    "        with torch.no_grad(), torch.cuda.stream(stream):\n",
    "            torch_inputs = torch.from_numpy(features).to(device)\n",
    "            outputs = trt_gm(torch_inputs)\n",
    "            return {\n",
    "                \"preds\": outputs.cpu().numpy(),\n",
    "            }\n",
    "\n",
    "    workspace_path = f\"/tmp/triton_{time.strftime('%m_%d_%M_%S')}\"\n",
    "    triton_conf = TritonConfig(http_port=ports[0], grpc_port=ports[1], metrics_port=ports[2])\n",
    "    with Triton(config=triton_conf, workspace=workspace_path) as triton:\n",
    "        triton.bind(\n",
    "            model_name=\"HousingModel\",\n",
    "            infer_func=_infer_fn,\n",
    "            inputs=[\n",
    "                Tensor(name=\"features\", dtype=np.float32, shape=(-1,)),\n",
    "            ],\n",
    "            outputs=[\n",
    "                Tensor(name=\"preds\", dtype=np.float32, shape=(-1,)),\n",
    "            ],\n",
    "            config=ModelConfig(\n",
    "                max_batch_size=50,\n",
    "                batcher=DynamicBatcher(max_queue_delay_microseconds=5000),  # 5ms\n",
    "            ),\n",
    "            strict=True,\n",
    "        )\n",
    "\n",
    "        def _stop_triton(signum, frame):\n",
    "            # The server manager sends SIGTERM to stop the server; this function ensures graceful cleanup.\n",
    "            print(\"SERVER: Received SIGTERM. Stopping Triton server.\")\n",
    "            triton.stop()\n",
    "\n",
    "        signal.signal(signal.SIGTERM, _stop_triton)\n",
    "\n",
    "        print(\"SERVER: Serving inference\")\n",
    "        triton.serve()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74121cd7",
   "metadata": {},
   "source": [
    "#### Start Triton servers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d6b7143",
   "metadata": {},
   "source": [
    "The `TritonServerManager` will handle the lifecycle of Triton server instances across the Spark cluster:\n",
    "- Find available ports for HTTP/gRPC/metrics\n",
    "- Deploy a server on each node via stage-level scheduling\n",
    "- Gracefully shutdown servers across nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "2fb22db8",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"HousingModel\"\n",
    "server_manager = TritonServerManager(model_name=model_name, model_path=model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e067aa14",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-07 11:03:44,809 - INFO - Requesting stage-level resources: (cores=5, gpu=1.0)\n",
      "2025-02-07 11:03:44,810 - INFO - Starting 1 servers.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'cb4ae00-lcedt': (2020631, [7000, 7001, 7002])}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Returns {'hostname', (server_pid, [http_port, grpc_port, metrics_port])}\n",
    "server_manager.start_servers(triton_server)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a1ac038",
   "metadata": {},
   "source": [
    "#### Define client function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4ac45ef",
   "metadata": {},
   "source": [
    "Get the hostname -> url mapping from the server manager:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92760dac",
   "metadata": {},
   "outputs": [],
   "source": [
    "host_to_http_url = server_manager.host_to_http_url  # or server_manager.host_to_grpc_url"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "122ebe7c",
   "metadata": {},
   "source": [
    "Define the Triton inference function, which returns a predict function for batch inference through the server:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "1ae91c54",
   "metadata": {},
   "outputs": [],
   "source": [
    "def triton_fn(model_name, host_to_url):\n",
    "    import socket\n",
    "    from pytriton.client import ModelClient\n",
    "\n",
    "    url = host_to_url[socket.gethostname()]\n",
    "    print(f\"Connecting to Triton model {model_name} at {url}.\")\n",
    "\n",
    "    def infer_batch(inputs):\n",
    "        with ModelClient(url, model_name, inference_timeout_s=240) as client:\n",
    "            result_data = client.infer_batch(inputs)\n",
    "            return result_data[\"preds\"]\n",
    "        \n",
    "    return infer_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "d3e64fda-117b-4810-a9a2-dd498239496f",
   "metadata": {},
   "outputs": [],
   "source": [
    "regress = predict_batch_udf(partial(triton_fn, model_name=model_name, host_to_url=host_to_http_url),\n",
    "                               input_tensor_shapes=[[8]],\n",
    "                               return_type=FloatType(),\n",
    "                               batch_size=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20b8514e-01de-481f-86aa-75afd99bcc7c",
   "metadata": {},
   "source": [
    "### Run Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "5eae04bc-75ca-421a-87c8-ac507ce1f2f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = spark.read.parquet(data_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "b350bd8e-9b8f-4511-9ddf-76d917b21b5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = df.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "a24149a5-3adc-4089-8769-13cf1e44547a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 16:>                                                         (0 + 8) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 25.8 ms, sys: 6.21 ms, total: 32.1 ms\n",
      "Wall time: 2.37 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# first pass caches model/fn\n",
    "predictions = df.withColumn(\"preds\", regress(struct(*columns)))\n",
    "preds = predictions.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "df2ce39f-30af-491a-8472-800fb1ce8458",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 17:>                                                         (0 + 8) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 171 ms, sys: 3.76 ms, total: 174 ms\n",
      "Wall time: 2.5 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "predictions = df.withColumn(\"preds\", regress(array(*columns)))\n",
    "preds = predictions.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "ca6f3eaa-9569-45d0-88bf-9aa0757e1ecb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 18:>                                                         (0 + 8) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 24.4 ms, sys: 4.83 ms, total: 29.2 ms\n",
      "Wall time: 1.97 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "predictions = df.withColumn(\"preds\", regress(array(*columns)))\n",
    "preds = predictions.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "b79c62c8-e1e8-4467-8aef-8939c31833b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------------+------------+-----------+------------+-----------+------------+----------+------------+---------+\n",
      "|      MedInc|    HouseAge|   AveRooms|   AveBedrms| Population|    AveOccup|  Latitude|   Longitude|    preds|\n",
      "+------------+------------+-----------+------------+-----------+------------+----------+------------+---------+\n",
      "|  0.20909257|  -1.1632254| 0.38946992|  0.04609274| -0.9806099| -0.07099328|0.61245227|-0.020113053|1.3746364|\n",
      "|-0.098627955|  0.34647804| 0.27216315|  -0.0129226| -0.6953838| -0.05380849| 1.0665938|  -1.2479742|1.8087528|\n",
      "| -0.66006273|   1.0616008|-0.55292207| -0.48945764|-0.13641118| 0.028952759| 1.1040496|  -1.3827378|1.4245079|\n",
      "|  0.08218294|   0.5848523|-0.13912922| -0.14707813|-0.19116047| -0.07136432|0.96827507|  -1.3028787|2.3895802|\n",
      "|   0.0784456|  -1.4810578| 0.57265776|  0.32067496|  1.0345173|-0.024157424| 1.4411427| -0.52423614|1.3616933|\n",
      "| -0.82318723| -0.36864465| 0.07829511|  -0.1808107|-0.67242444|-0.061470542| 1.9374212|  -1.0083897|0.7539238|\n",
      "|  0.59671736|   0.5848523| 0.19346413|  -0.1371872|-0.19645879| 0.009964322|0.96827507|  -1.2928978|2.6816423|\n",
      "|  -0.9612035|  -1.5605159|-0.56329846| 0.027148023|-0.71127874| -0.08471591| 0.5328614| -0.13990337|1.1731354|\n",
      "| -0.74344087|  -1.2426835| 0.27282518|   0.4037246| -0.9841421| -0.05610115| 1.2257773| -0.42940006|1.0198532|\n",
      "|   0.9784464|  -0.2891866| 0.24374022| -0.24670053| 0.28922042| -0.01102468| 1.1087307|  -1.2280084| 2.708211|\n",
      "|  -0.5070446|  -1.0043093|-0.78254056|0.0122275995|  2.8465424|-0.060435444| 0.8980464|  -1.2080427|2.0327075|\n",
      "| -0.18690155|   1.2205169|0.015323491|  0.12183313|-0.41015765|  0.04452552|  1.010412|  -1.3228445|1.9909104|\n",
      "|  -1.2551856|   1.6178073| -0.3341509|-0.060125165| -0.7554314| -0.08777025| 1.0291398|  -1.3477987|1.2702764|\n",
      "|   4.9607058|  -1.9578062|  1.4854684| -0.03948475|  2.1833694|0.0029250523|  1.024457|  -1.1581304| 5.975229|\n",
      "|  0.73652315|  -1.6399739|  0.7913185| -0.05238397|    1.67738|  0.01944797| 1.0993668|  -1.1331724|1.9309721|\n",
      "|   -0.505834|  0.18756187|-0.47093546| -0.24297306|-0.60619545| -0.10791535|  0.977639|  -1.2879055|1.7610806|\n",
      "| -0.88477343|-0.050812364| -0.6318951| -0.15244243| -0.5258376| -0.15618815| 0.9823201|  -1.2879055| 1.655031|\n",
      "| -0.42840376|   0.9821427| -0.2266495| -0.36083496| -0.6883194| -0.08552282| 0.5328614| -0.12493005|1.1175063|\n",
      "|   0.9369153|  -1.4810578|  0.6722208|-0.121177554|  0.3996021|  0.01291408| 1.1040496|  -1.1082181|2.1779811|\n",
      "| -0.80702734| -0.92485124|-0.26602685|  -0.1560743|  1.4398388| -0.09314839|0.55627036| -0.09498342|0.9102398|\n",
      "+------------+------------+-----------+------------+-----------+------------+----------+------------+---------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "predictions.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fec23b0-eaf2-4b6a-aa38-7a09873ed6eb",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Stop Triton Server on each executor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "8084bdef",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[True]"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "server_manager.stop_servers()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "0138a029-87c5-497f-ac5c-3eed0e11b0f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not on_databricks: # on databricks, spark.stop() puts the cluster in a bad state\n",
    "    spark.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d24147e7-5695-44a0-9961-b94bfba1cfff",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "spark-dl-torch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
